# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：SC_half_plane_real.py
    Created on：2025/6/5 22:33
    Author：(input)
    Description:
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.tri as tri
from scipy.interpolate import griddata

class BodyFittedGrid:
    """
    Body-Fitted Coordinates Grid
    """
    def __init__(self, L, h, alpha_deg):
        self.L = L
        self.h = h
        self.alpha = np.deg2rad(alpha_deg)

        # Physical domain points (corrected order: [0,0], [L,0], [L,h], [-h*tan(α),h])
        A1 = [0.0, 0.0]  # (xi=0, eta=0)
        A2 = [L, 0.0]  # (xi=1, eta=0)
        A3 = [L, h]  # (xi=1, eta=1)
        A4 = [-h * np.tan(self.alpha), h]  # (xi=0, eta=1)
        self.P = np.array([A1, A2, A3, A4])

        # Computational domain points (aligned with physical order)
        self.Lp = L + h * np.tan(self.alpha)
        self.hp = h
        A1p = [0.0, 0.0]  # (xi=0, eta=0)
        A2p = [self.Lp, 0.0]  # (xi=1, eta=0)
        A3p = [self.Lp, self.hp]  # (xi=1, eta=1)
        A4p = [0.0, self.hp]  # (xi=0, eta=1)
        self.Q = np.array([A1p, A2p, A3p, A4p])

    def shape_functions(self, xi, eta):
        N0 = (1 - xi) * (1 - eta)  # (0,0)
        N1 = xi * (1 - eta)  # (1,0)
        N2 = xi * eta  # (1,1)
        N3 = (1 - xi) * eta  # (0,1)
        return np.array([N0, N1, N2, N3])

    def map_to_physical(self, xi, eta):
        N = self.shape_functions(xi, eta)
        x = np.dot(N, self.P[:, 0])
        y = np.dot(N, self.P[:, 1])
        return x, y

    def map_to_computational(self, xi, eta):
        N = self.shape_functions(xi, eta)
        x = np.dot(N, self.Q[:, 0])
        y = np.dot(N, self.Q[:, 1])
        return x, y

    def map_to_upper_half(self, xi, eta):
        x, y = self.map_to_computational(xi, eta)
        z = x + 1j * y
        w = np.cosh(np.pi / self.hp * z)
        return w

    def inverse_map_physical_to_xieta(self, x_phys, y_phys, tol=1e-8, max_iter=100):
        xi, eta = 0.5, 0.5
        for _ in range(max_iter):
            x_cur, y_cur = self.map_to_physical(xi, eta)
            dx_dxi = (self.map_to_physical(xi + tol, eta)[0] - self.map_to_physical(xi - tol, eta)[0]) / (2 * tol)
            dx_deta = (self.map_to_physical(xi, eta + tol)[0] - self.map_to_physical(xi, eta - tol)[0]) / (2 * tol)
            dy_dxi = (self.map_to_physical(xi + tol, eta)[1] - self.map_to_physical(xi - tol, eta)[1]) / (2 * tol)
            dy_deta = (self.map_to_physical(xi, eta + tol)[1] - self.map_to_physical(xi, eta - tol)[1]) / (2 * tol)
            J = np.array([[dx_dxi, dx_deta], [dy_dxi, dy_deta]])
            rhs = np.array([x_phys - x_cur, y_phys - y_cur])
            delta = np.linalg.solve(J, rhs)
            xi += delta[0]
            eta += delta[1]
            if np.linalg.norm(delta) < tol:
                break
        return xi, eta

    def generate_grid(self, Nx, Ny, target='physical'):
        xi_vals = np.linspace(0, 1, Nx)
        eta_vals = np.linspace(0, 1, Ny)
        grid = []

        for eta in eta_vals:
            row = []
            for xi in xi_vals:
                if target == 'physical':
                    x, y = self.map_to_physical(xi, eta)
                elif target == 'computational':
                    x, y = self.map_to_computational(xi, eta)
                elif target == 'upper_half':
                    w = self.map_to_upper_half(xi, eta)
                    x, y = w.real, w.imag
                row.append((x, y))
            grid.append(row)

        return np.array(grid)  # shape (Ny, Nx, 2)



class ComplexPotential:
    def __init__(self, mapper, source_phys_pos_list, Qc_list, margin=0.1):
        """
        Parameters:
        - mapper: 必须提供 map_to_physical(xi, eta), map_to_computational(xi, eta), map_to_upper_half(xi, eta) 方法
        - source_phys_pos_list: 源点在物理平面上的位置列表 [(x1, y1), (x2, y2), ...]
        - Qc_list: 源强列表 [Qc1, Qc2, ...]
        - margin: 绘图边距比例
        """
        self.mapper = mapper
        self.source_phys_list = [np.array(pos) for pos in source_phys_pos_list]
        self.Qc_list = Qc_list
        self.margin = margin
        self.h_prime = self.mapper.Q[:, 1].max()

        if len(self.source_phys_list) != len(self.Qc_list):
            raise ValueError("Number of source positions must match number of source strengths")

        self.xi_eta_sources = []
        self.wc_list = []
        for idx, source_phys in enumerate(self.source_phys_list):
            if not self._is_point_in_physical_domain(source_phys):
                raise ValueError(f"Source point {source_phys} at index {idx} is outside the physical domain")
            self.compute_source_position(source_phys)

    def _is_point_in_physical_domain(self, point):
        """
        使用射线法判断点是否在物理域平行四边形内，包括边界点。
        Parameters:
        - point: 点坐标 (x, y)
        Returns:
        - bool: 点是否在物理域内（包括边界）
        """
        x, y = point
        vertices = self.mapper.P  # [[0,0], [L,0], [L,h], [-h*tan(α),h]]

        # 检查是否直接位于顶点上
        for vertex in vertices:
            if np.allclose([x, y], vertex, atol=1e-10):
                return True

        # 射线法：从点向右发射射线，计算与边界的交点次数
        inside = False
        for i in range(len(vertices)):
            j = (i + 1) % len(vertices)
            xi, yi = vertices[i]
            xj, yj = vertices[j]

            # 包含边界点：使用 >= 和 <=
            if ((yi > y) != (yj > y)) or ((yi == y) and (yj == y)):
                # 计算交点 x 坐标
                if yi == yj:  # 水平边
                    if min(xi, xj) <= x <= max(xi, xj):
                        inside = not inside
                else:
                    x_intersect = (xj - xi) * (y - yi) / (yj - yi) + xi
                    if x_intersect > x:
                        inside = not inside

        return inside

    def compute_source_position(self, source_phys):
        """
        计算单个物理源点对应的 (xi, eta)，并映射到计算平面和复平面。
        """
        x0, y0 = source_phys
        h = self.mapper.h
        L = self.mapper.L
        tan_alpha = np.tan(self.mapper.alpha)

        eta = y0 / h
        xi = (x0 + y0 * tan_alpha) / (L + y0 * tan_alpha)

        xi_eta = (np.clip(xi, 0, 1), np.clip(eta, 0, 1))
        self.xi_eta_sources.append(xi_eta)

        wc = self.mapper.map_to_upper_half(*xi_eta)
        self.wc_list.append(wc)

    def compute_point_source_field(self, Nx=100, Ny=100):
        """
        在整个计算区域生成多个点源叠加的速度势、流函数与速度分量
        """
        xi_vals = np.linspace(0, 1, Nx)
        eta_vals = np.linspace(0, 1, Ny)

        data = []
        for eta in eta_vals:
            for xi in xi_vals:
                x_phys, y_phys = self.mapper.map_to_physical(xi, eta)
                w = self.mapper.map_to_upper_half(xi, eta)

                W_total = 0
                u_total = 0
                v_total = 0
                for Qc, wc in zip(self.Qc_list, self.wc_list):
                    W = (Qc / np.pi) * np.log(w - wc)
                    W_total += W
                    dW_dw = (Qc / np.pi) / (w - wc)
                    u_total += np.real(dW_dw)
                    v_total += np.imag(dW_dw)

                phi = np.real(W_total)
                psi = np.imag(W_total)
                u = u_total
                v = v_total
                magV = np.hypot(u, v)

                data.append({
                    "x": x_phys,
                    "y": y_phys,
                    "phi": phi,
                    "psi": psi,
                    "u": u,
                    "v": v,
                    "magV": magV
                })

        df = pd.DataFrame(data)
        return df

    def plot_boundary(self):
        '''
        绘制边界、源点位置
        '''
        x, y = df["x"], df["y"]
        # 2. 绘制物理边界（参考 plot_grids 的方式）
        edges = [(0, 3), (2, 3), (0, 1)]  # 左边界 (0->3), 顶边界 (3->2), 底边界 (0->1)
        for idx, (i, j) in enumerate(edges):
            x_vals = [self.mapper.P[i][0], self.mapper.P[j][0]]
            y_vals = [self.mapper.P[i][1], self.mapper.P[j][1]]
            label = {0: 'Inclined wall', 1: 'Top wall', 2: 'Bottom wall'}[idx]
            plt.plot(x_vals, y_vals, 'k-', linewidth=2.5, label=label)

        # 3. 蓝色参考线：x = 0
        y_ref = np.linspace(y.min(), y.max(), 200)
        x_ref = np.zeros_like(y_ref)
        plt.plot(x_ref, y_ref, 'k--', linewidth=1.8, label='x = 0 ref line')

        # 4. 点源标记（支持多个源点）
        for idx, (source_x, source_y) in enumerate(self.source_phys_list):
            plt.plot(source_x, source_y, 'ro', markersize=6, label='Source Point' if idx == 0 else "")

    def plot_streamlines(self, df, levels=30):
        """
        绘制流函数等值线图（streamlines），包括物理边界、参考线和源点。

        Parameters:
        - df: DataFrame 包含 x, y, psi 列
        - levels: 等值线数量 (默认 30)
        """
        x, y, psi, phi = df["x"], df["y"], df["psi"], df["phi"]
        margin_x = self.margin * (x.max() - x.min())
        margin_y = self.margin * (y.max() - y.min())

        triang = tri.Triangulation(x, y)
        plt.figure(figsize=(7, 6))

        cv = plt.tricontour(triang, phi, levels=levels, colors='r', linewidths=0.6)
        plt.clabel(cv, inline=True, fontsize=8, fmt="%.1f")


        # 1. 绘制流线
        cs = plt.tricontour(triang, psi, levels=levels, colors='b', linewidths=0.6)
        plt.clabel(cs, inline=True, fontsize=8, fmt="%.1f")

        # 234 绘制边界点源
        self.plot_boundary()

        # 5. 图形设置
        plt.title("Streamlines (ψ) with Physical Boundaries")
        plt.xlim(x.min() - margin_x, x.max() + margin_x)
        plt.ylim(y.min() - margin_y, y.max() + margin_y)
        plt.gca().set_aspect('equal')
        plt.xlabel("x")
        plt.ylabel("y")
        plt.grid(True)
        #plt.legend(loc="best")
        plt.tight_layout()
        plt.show()

    def plot_velocity_contours(self, df, levels=50, show_boundaries=True,
                               title='Velocity Magnitude Contours', cmap='jet'):
        """
        使用原始非规则网格数据绘制速度模 V 的等值线图（不插值）

        Parameters
        ----------
        df : DataFrame
            包含 'x', 'y', 'magV' 的数据
        levels : int
            等值线层数
        show_boundaries : bool
            是否绘制物理边界
        title : str
            图标题
        cmap : str
            色图名称（默认 'jet'）
        """
        x = df["x"].values
        y = df["y"].values
        V = df["magV"].values

        # 设置边距范围（同流函数绘图）
        margin_x = self.margin * (x.max() - x.min())
        margin_y = self.margin * (y.max() - y.min())
        x_min, x_max = x.min() - margin_x, x.max() + margin_x
        y_min, y_max = y.min() - margin_y, y.max() + margin_y

        # 三角剖分网格
        triang = tri.Triangulation(x, y)

        # 绘图
        plt.figure(figsize=self.figsize if hasattr(self, 'figsize') else (6, 5))
        contour = plt.tricontourf(triang, V, levels=levels, cmap=cmap)
        plt.clabel(contour, inline=True, fontsize=8, fmt="%.1f")
        plt.colorbar(contour, label="Velocity Magnitude")

        # 边界线
        if show_boundaries:
            self.plot_boundary()

        plt.title(title)
        plt.xlabel("x")
        plt.ylabel("y")
        plt.xlim(x_min, x_max)
        plt.ylim(y_min, y_max)
        plt.gca().set_aspect('equal')
        plt.grid(True, linestyle='--', alpha=0.4)
        plt.tight_layout()
        plt.show()

    def plot_quantity_along_line(self, df, point_start, point_end, variable='V', N=200,
                                 normalize=False, label=None, title=None):
        """
        沿任意指定线段绘制物理量分布图

        Parameters
        ----------
        df : DataFrame
            包含 'x', 'y', 和指定物理量列（如 'V', 'phi', 'psi'）的流场数据
        point_start : tuple
            起点坐标 (x0, y0)
        point_end : tuple
            终点坐标 (x1, y1)
        variable : str
            要提取和绘制的物理量列名（如 'V', 'phi', 'u'）
        N : int
            插值点数
        normalize : bool
            是否对物理量进行无量纲归一化（/ max）
        label : str or None
            曲线图例名称
        title : str or None
            图标题
        """
        from scipy.interpolate import griddata

        # 插值目标点线段
        x0, y0 = point_start
        x1, y1 = point_end
        t = np.linspace(0, 1, N)
        x_line = x0 + (x1 - x0) * t
        y_line = y0 + (y1 - y0) * t

        # 原始数据
        x_data = df["x"].values
        y_data = df["y"].values
        val_data = df[variable].values

        # 插值物理量
        val_line = griddata((x_data, y_data), val_data, (x_line, y_line), method='linear')

        if normalize and np.max(np.abs(val_line)) > 1e-8:
            val_line /= np.max(np.abs(val_line))

        # 绘图
        plt.figure(figsize=(6, 4))
        plt.plot(t, val_line, 'b-', lw=2, label=label or f"{variable} profile")

        plt.xlabel("Normalized Distance Along Line (t)")
        plt.ylabel(variable + (" / max" if normalize else ""))
        plt.title(title or f"{variable} along line [{point_start} → {point_end}]")
        plt.grid(True)
        if label:
            plt.legend()
        plt.tight_layout()
        plt.show()

    def compute_quantity_along_line(self, point_start, point_end, variable='V', N=200, normalize=False, title=None):
        """
        沿给定线段计算并绘制指定物理量的分布

        Parameters
        ----------
        point_start : tuple (x0, y0)
            起点坐标
        point_end : tuple (x1, y1)
            终点坐标
        variable : str
            要计算的物理量名称：'phi', 'psi', 'u', 'v', 'V'
        N : int
            线段上的点数
        normalize : bool
            是否进行最大值归一化
        title : str
            图标题（默认自动生成）
        """

        x0, y0 = point_start
        x1, y1 = point_end
        t = np.linspace(0, 1, N)
        x_line = x0 + (x1 - x0) * t
        y_line = y0 + (y1 - y0) * t

        # 初始化数据
        values = []

        for x, y in zip(x_line, y_line):
            # 将物理坐标 (x,y) 映射为 (xi, eta)
            xi, eta = self.mapper.inverse_map_physical_to_xieta(x, y)

            # 映射到辅助平面 w
            w = self.mapper.map_to_upper_half(xi, eta)

            # 复势及导数
            W = (self.Qc_list[0] / np.pi) * np.log(w - self.wc_list[0])
            dW_dw = (self.Qc_list[0] / np.pi) / (w - self.wc_list[0])

            # 选择物理量
            if variable == 'phi':
                val = np.real(W)
            elif variable == 'psi':
                val = np.imag(W)
            elif variable == 'u':
                val = np.real(dW_dw)
            elif variable == 'v':
                val = np.imag(dW_dw)
            elif variable == 'V':
                val = np.abs(dW_dw)
            else:
                raise ValueError(f"Unsupported variable: {variable}")

            values.append(val)

        values = np.array(values)

        # 归一化
        if normalize and np.max(np.abs(values)) > 1e-6:
            values /= np.max(np.abs(values))

        # 绘图
        plt.figure(figsize=(6, 4))
        plt.plot(t, values, 'b-', lw=2)
        plt.xlabel("Normalized distance (t)")
        plt.ylabel(f"{variable} {'(normalized)' if normalize else ''}")
        plt.title(title or f"{variable} along line [{point_start} → {point_end}]")
        plt.grid(True)
        plt.tight_layout()
        plt.show()


def plot_grids(grid, Nx, Ny, target='physical', source_phys_pos=None, title_prefix="Grid"):
    """
    绘制物理平面、计算平面或上半复平面的单个网格图，包含网格点、边界线和源点（仅物理平面）。

    参数:
        grid: BodyFittedGrid 实例
        Nx, Ny: xi 和 eta 方向的网格点数
        target: 绘制平面，选项为 'physical', 'computational', 'upper_half'
        source_phys_pos: 元组 (x, y)，物理平面上的源点坐标（仅对 physical 有效）
        title_prefix: 图形标题前缀
    """
    # 生成指定平面的网格
    grid_data = grid.generate_grid(Nx, Ny, target=target)
    X, Y = grid_data[:, :, 0], grid_data[:, :, 1]

    # 创建单个图形
    plt.figure(figsize=(6, 6))

    # 绘制网格点
    plt.scatter(X, Y, color='red', s=5, label='Grid Nodes')

    # 边界线索引（连接 P[0]-P[3], P[2]-P[3], P[0]-P[1]，左边界和上下边界）
    edges = [(0, 3), (2, 3), (0, 1)]#

    if target == 'physical':
        # 绘制参考线 [0, 0] 到 [0, h]
        plt.plot([0, 0], [0, grid.h], 'k--', linewidth=1, label='Reference Line')
        # 绘制边界线
        for i, j in edges:
            x_vals = [grid.P[i][0], grid.P[j][0]]
            y_vals = [grid.P[i][1], grid.P[j][1]]
            plt.plot(x_vals, y_vals, 'b-', linewidth=2, label='Boundary' if i == 0 else None)
        # 绘制源点
        if source_phys_pos is not None:
            x0, y0 = source_phys_pos
            plt.scatter(x0, y0, color='green', marker='o', s=20, label='Source Point')
        plt.title(f'{title_prefix} - Physical Domain')
        plt.xlabel('x')
        plt.ylabel('y')

    elif target == 'computational':
        # 绘制边界线
        for i, j in edges:
            x_vals = [grid.Q[i][0], grid.Q[j][0]]
            y_vals = [grid.Q[i][1], grid.Q[j][1]]
            plt.plot(x_vals, y_vals, 'b-', linewidth=2, label='Boundary' if i == 0 else None)
        plt.title(f'{title_prefix} - Computational Domain')
        plt.xlabel('x')
        plt.ylabel('y')

    elif target == 'upper_half':
        # 绘制网格线（代替边界）
        for row in grid_data:
            plt.plot(row[:, 0], row[:, 1], 'b-', alpha=0.5, label='ξ lines' if row is grid_data[0] else "")
        for col in grid_data.transpose(1, 0, 2):
            plt.plot(col[:, 0], col[:, 1], 'r-', alpha=0.2, label='η lines' if col is grid_data[:, 0] else "")
        plt.title(f'{title_prefix} - Upper Half Plane')
        plt.xlabel('Re(w)')
        plt.ylabel('Im(w)')

    else:
        raise ValueError("target must be 'physical', 'computational', or 'upper_half'")

    plt.gca().set_aspect('equal')
    plt.grid(True, linestyle='--', alpha=0.3)
    #plt.legend()
    plt.show()

# 可视化（支持多个源点）
def plot_grids_with_multiple_sources(grid, Nx, Ny, source_phys_pos_list, target='physical', title_prefix="Grid"):
    grid_data = grid.generate_grid(Nx, Ny, target=target)
    X, Y = grid_data[:, :, 0], grid_data[:, :, 1]

    plt.figure(figsize=(6, 6))
    plt.scatter(X, Y, color='red', s=5, label='Grid Nodes')

    edges = [(0, 3), (2, 3), (0, 1)]
    if target == 'physical':
        plt.plot([0, 0], [0, grid.h], 'k--', linewidth=1, label='Reference Line')
        for i, j in edges:
            x_vals = [grid.P[i][0], grid.P[j][0]]
            y_vals = [grid.P[i][1], grid.P[j][1]]
            plt.plot(x_vals, y_vals, 'b-', linewidth=2, label='Boundary' if i == 0 else None)
        for idx, source_pos in enumerate(source_phys_pos_list):
            x0, y0 = source_pos
            plt.scatter(x0, y0, color='green', marker='X', s=20, label=f'Source Point {idx+1}')
    # 其他平面逻辑保持不变（省略）

    plt.gca().set_aspect('equal')
    plt.grid(True, linestyle='--', alpha=0.3)
    #plt.legend()
    plt.show()



if __name__ == "__main__":
    # 参数设置
    L, h, alpha_deg = 10.0, 2.0, -10
    s = 0.6
    U0 = 20
    b = 0.02
    source_phys_pos = (s, h)  # 单源点用于绘图
    Qc = U0 * b
    source_phys_pos_list = [(0.6, 2.0)]  # 多个源点
    Qc_list = [Qc]

    # 创建网格实例
    grid = BodyFittedGrid(L, h, alpha_deg)

    # 测试一组 (xi, eta) 值
    xi = 0.5
    eta = 0.5
    w = grid.map_to_upper_half(xi, eta)
    print(f"map_to_upper_half 输出: {w}")

    # 绘制网格
    Nx, Ny = 20, 20  # 网格分辨率
    #plot_grids(grid, Nx, Ny, target='physical', source_phys_pos=source_phys_pos, title_prefix="Grid Visualization")
    # plot_grids(grid, Nx, Ny, target='computational', title_prefix="Grid Visualization")
    # plot_grids(grid, Nx, Ny, target='upper_half', title_prefix="Grid Visualization")

    # 创建势场实例
    potential = ComplexPotential(grid, source_phys_pos_list, Qc_list)

    # 验证多个源点位置
    print(f"xi_eta_sources: {potential.xi_eta_sources}")
    for idx, (xi, eta) in enumerate(potential.xi_eta_sources):
        print(f"Source {idx} (xi, eta): ({xi:.4f}, {eta:.4f})")
        x_phys, y_phys = grid.map_to_physical(xi, eta)
        print(f"Mapped back to physical: ({x_phys:.4f}, {y_phys:.4f})")
        print(f"Expected physical: ({source_phys_pos_list[idx][0]:.4f}, {source_phys_pos_list[idx][1]:.4f})")

    # 计算场
    df = potential.compute_point_source_field(Nx=50, Ny=50)
    print("Field data sample:")
    print(df.head())

    # 绘制网格和多个源点
    #plot_grids_with_multiple_sources(grid, Nx, Ny, source_phys_pos_list, target='physical', title_prefix="Grid with Sources")

    # 绘制流线
    potential.plot_streamlines(df, levels=10)

    # 绘制速度等值线
    #potential.plot_velocity_contours(df)

    # 提取沿 x = 0.5 垂直线，y = [0, h] 上的 V 分布
    #potential.plot_quantity_along_line(df, point_start=(0.5, 0), point_end=(0.5, h), variable='magV', normalize=True)

    # 提取沿斜壁面（A3 → A2）上的 φ 分布
    #potential.plot_quantity_along_line(df, point_start=(0.0, 0.0), point_end=(-h * np.tan(alpha_deg), h), variable='phi')

    # 沿 x = 0.5 的垂直线段，计算 V 分布
    potential.compute_quantity_along_line((0.5, 0.0), (0.5, h), variable='V', normalize=True)

    # 沿倾斜壁面 A3 → A2，计算 ψ 分布
    A2 = [-h * np.tan(alpha_deg), h]
    potential.compute_quantity_along_line((0.0, 0.0), A2, variable='psi')

